# app.py
# -*- coding: utf-8 -*-
"""
Flask 工作流 Demo（模板版）
- 提供：工作流 REST 接口、SSE 订阅接口、首页模板渲染
- 优先复用 wiring.container 中的 orchestrator / eventbus；无则用内置简化实现
运行：
  pip install flask
  python app.py
  打开 http://127.0.0.1:5000
"""

from __future__ import annotations

import json
import time
import queue
import asyncio
from typing import Any, Dict, Optional

from flask import Flask, request, Response, jsonify, render_template

# ============================================================
# 1) 依赖装配：优先尝试导入你项目的容器；失败则启用最小内置实现
# ============================================================
try:
    # ✅ 如果你有正式的容器，直接复用它（强烈推荐）
    from wiring.container import orchestrator, eventbus  # 改成你项目的导入路径
    print("[BOOT] use container.orchestrator / container.eventbus")
except Exception:
    print("[BOOT] container not found, fallback to minimal in-memory impl.")

    # ---- 最小内置实现（仅为本地 Demo）----
    # 内存事件总线（queue.Queue 订阅/发布）
    class MemoryEventBus:
        def __init__(self, max_queue_size: int = 1000):
            self._subs: Dict[str, list[queue.Queue]] = {}
            self._maxsize = max_queue_size

        def create_queue(self) -> queue.Queue:
            return queue.Queue(maxsize=self._maxsize)

        def subscribe(self, exec_id: str, q: queue.Queue) -> None:
            self._subs.setdefault(exec_id, []).append(q)
            print(f"[BUS] subscribe exec={exec_id} subs={len(self._subs[exec_id])}")

        def unsubscribe(self, exec_id: str, q: queue.Queue) -> None:
            lst = self._subs.get(exec_id, [])
            if q in lst:
                lst.remove(q)
            if not lst and exec_id in self._subs:
                self._subs.pop(exec_id, None)
            print(f"[BUS] unsubscribe exec={exec_id}")

        def publish(self, exec_id: str, event: str, payload: dict) -> None:
            for q in list(self._subs.get(exec_id, [])):
                try:
                    q.put_nowait({"event": event, "payload": payload})
                except queue.Full:
                    # 丢弃最老一条再塞
                    try:
                        _ = q.get_nowait()
                    except queue.Empty:
                        pass
                    try:
                        q.put_nowait({"event": event, "payload": payload})
                    except queue.Full:
                        pass

    eventbus = MemoryEventBus()

    # 枚举（简化）
    class EventType:
        EXEC_STARTED   = "exec_started"
        EXEC_FINISHED  = "exec_finished"
        EXEC_PAUSED    = "exec_paused"
        EXEC_RESUMED   = "exec_resumed"
        EXEC_CANCELLING= "exec_cancelling"
        NODE_STARTED   = "node_started"
        NODE_FINISHED  = "node_finished"
        NODE_FAILED    = "node_failed"
        NODE_CANCELLED = "node_cancelled"

    class ExecStatus:
        PENDING   = "PENDING"
        RUNNING   = "RUNNING"
        SUCCEEDED = "SUCCEEDED"
        FAILED    = "FAILED"
        CANCELLED = "CANCELLED"

    # 状态存储（极简）
    import uuid
    class MemoryState:
        def __init__(self):
            self._store: Dict[str, dict] = {}
        def init(self, eid: str):
            self._store[eid] = {
                "status": ExecStatus.PENDING,
                "running": [],
                "finished": [],
                "errors": {},
                "created_at": time.time(),
                "updated_at": time.time(),
                "finished_at": None,
            }
        def update(self, eid: str, **kw):
            if eid in self._store:
                self._store[eid].update(kw)
                self._store[eid]["updated_at"] = time.time()
        def snapshot(self, eid: str) -> dict:
            return dict(self._store[eid])
    state = MemoryState()

    # 最小编排服务：启动一个异步任务模拟执行（3秒 code 节点）
    class MiniOrchestrator:
        def __init__(self):
            self._tasks: Dict[str, asyncio.Task] = {}
            self._loop = asyncio.get_event_loop()

        def start(self, workflow: Dict, context: Optional[Dict] = None, exec_id: Optional[str] = None) -> str:
            eid = exec_id or str(uuid.uuid4())
            state.init(eid)
            state.update(eid, status=ExecStatus.RUNNING)
            eventbus.publish(eid, EventType.EXEC_STARTED, {"exec_id": eid, "ready": ["code_1"]})
            # 启后台任务
            async def _run():
                try:
                    eventbus.publish(eid, EventType.NODE_STARTED, {"exec_id": eid, "node_id": "code_1"})
                    await asyncio.sleep(3)
                    eventbus.publish(eid, EventType.NODE_FINISHED, {"exec_id": eid, "node_id": "code_1", "result": {"ok": True}})
                    state.update(eid, status=ExecStatus.SUCCEEDED, finished_at=time.time())
                    eventbus.publish(eid, EventType.EXEC_FINISHED, {"exec_id": eid, "status": ExecStatus.SUCCEEDED})
                except asyncio.CancelledError:
                    state.update(eid, status=ExecStatus.CANCELLED, finished_at=time.time())
                    eventbus.publish(eid, EventType.EXEC_FINISHED, {"exec_id": eid, "status": ExecStatus.CANCELLED})
            self._tasks[eid] = self._loop.create_task(_run())
            return eid

        def snapshot(self, eid: str) -> dict:
            return state.snapshot(eid)

        def cancel(self, eid: str) -> None:
            t = self._tasks.get(eid)
            if t and not t.done():
                t.cancel()

        def pause(self, eid: str) -> None: pass
        def resume(self, eid: str) -> None: pass
        def cancel_node(self, eid: str, node_id: str) -> None: pass
        def list_active(self) -> Dict[str, str]:
            return {eid: state.snapshot(eid)["status"] for eid in self._tasks.keys()}

    orchestrator = MiniOrchestrator()

# ============================================================
# 2) Flask 应用
# ============================================================
app = Flask(__name__)

def ok(data: Any = None):
    return jsonify({"code": 200, "success": True, "message": "success", "result": data})

def err(msg: str, code: int = 400):
    return jsonify({"code": code, "success": False, "message": msg, "result": None}), code

# ---------- 页面 ----------
@app.get("/")
def index():
    return render_template("index.html")

# ---------- 工作流接口 ----------
@app.post("/executions/execute")
def start_execution():
    try:
        body = request.get_json(force=True) or {}
        wf   = body.get("workflow") or {}
        ctx  = body.get("context") or {}
        eid  = body.get("exec_id")
        exec_id = orchestrator.start(workflow=wf, context=ctx, exec_id=eid)
        return ok({"exec_id": exec_id, "status": "PENDING"})
    except Exception as e:
        return err(f"启动失败：{e}", 500)

@app.get("/executions/<exec_id>/status")
def get_status(exec_id: str):
    try:
        snap = orchestrator.snapshot(exec_id)
        return ok({"exec_id": exec_id, **snap})
    except Exception as e:
        return err(f"未找到执行：{e}", 404)

@app.post("/executions/<exec_id>/cancel")
def cancel(exec_id: str):
    try:
        orchestrator.cancel(exec_id)
        return ok(True)
    except Exception as e:
        return err(f"取消失败：{e}", 404)

@app.post("/executions/<exec_id>/pause")
def pause(exec_id: str):
    try:
        orchestrator.pause(exec_id)
        return ok(True)
    except Exception as e:
        return err(f"暂停失败：{e}", 404)

@app.post("/executions/<exec_id>/resume")
def resume(exec_id: str):
    try:
        orchestrator.resume(exec_id)
        return ok(True)
    except Exception as e:
        return err(f"恢复失败：{e}", 404)

@app.post("/executions/<exec_id>/nodes/<node_id>/cancel")
def cancel_node(exec_id: str, node_id: str):
    try:
        orchestrator.cancel_node(exec_id, node_id)
        return ok(True)
    except Exception as e:
        return err(f"取消节点失败：{e}", 404)

@app.get("/executions/active")
def list_active():
    items = orchestrator.list_active()
    out = [{"exec_id": k, "status": v} for k, v in items.items()]
    return ok(out)

# ---------- SSE 订阅 ----------
HEARTBEAT_SEC = 20
RETRY_MS = 3000

def _sse_bytes(event: str, data: dict) -> bytes:
    return f"event: {event}\n" f"data: {json.dumps(data, ensure_ascii=False)}\n\n".encode("utf-8")

def sse_stream(exec_id: str):
    q: queue.Queue = eventbus.create_queue()
    eventbus.subscribe(exec_id, q)
    try:
        # 首包
        yield f"retry: {RETRY_MS}\n\n".encode("utf-8")
        yield _sse_bytes("connected", {"exec_id": exec_id, "bus_id": id(eventbus), "ts": int(time.time())})

        last = time.time()
        while True:
            try:
                rec = q.get(timeout=1.0)
                yield _sse_bytes(rec.get("event", "message"), rec.get("payload", {}))
            except queue.Empty:
                pass

            now = time.time()
            if now - last >= HEARTBEAT_SEC:
                yield f": keepalive {int(now)}\n\n".encode("utf-8")
                last = now
    finally:
        eventbus.unsubscribe(exec_id, q)

@app.get("/executions/<exec_id>/subscribe")
def subscribe(exec_id: str):
    headers = {
        "Content-Type": "text/event-stream; charset=utf-8",
        "Cache-Control": "no-cache, no-transform",
        "X-Accel-Buffering": "no",
        "Connection": "keep-alive",
        "Content-Encoding": "identity",
    }
    return Response(sse_stream(exec_id), headers=headers, direct_passthrough=True)

# ---------- 一个固定的最小工作流 ----------
@app.post("/demo/start")
def demo_start():
    wf = {
        "nodes": [
            {"id": "start_0", "type": "start", "data": {"outputs": {"x": 1}}},
            {"id": "code_1", "type": "code", "data": {"codes": {"type": "python", "value": "async def main(args: Args) -> Output:\n    import asyncio\n    await asyncio.sleep(3)\n    p = args.params\n    return {\"result\": p.get(\"x\",0) + 1}\n"}}},
            {"id": "end_0", "type": "end", "data": {}},
        ],
        "edges": [
            {"sourceNodeID": "start_0", "targetNodeID": "code_1"},
            {"sourceNodeID": "code_1", "targetNodeID": "end_0"},
        ],
    }
    exec_id = orchestrator.start(workflow=wf, context={"y": 100})
    return ok({"exec_id": exec_id})

if __name__ == "__main__":
    # threaded=True 允许一个线程跑 SSE、另一个线程处理 REST 请求
    app.run(host="127.0.0.1", port=5000, debug=True, threaded=True)
